#/usr/bin/env/python
#resampler.py

"""Code for computing significance with jackknifing techniques.

Revision History:

5-23-05: Cathy Lozupone: started writing code
"""

from __future__ import division
from copy import deepcopy
from old_cogent.cluster.UPGMA import inputs_from_dict2D, UPGMA_cluster
from random import randrange
from old_cogent.base.stats import Numbers
from math import sqrt
from old_cogent.base.dict2d import Dict2D
from old_cogent.base.stats import Numbers
from old_cogent.cluster.metric_scaling import principal_coordinates_analysis
from Numeric import argsort, zeros, Float
from old_cogent.maths.stats.test import correlation
import pickle

BIG_NUM = 1e305

def median(x):
    """returns the median of list x"""
    #find the median of the low values
    x.sort()
    odd = len(x) % 2
    midpoint = int(len(x)/2)
    if odd:
        return x[midpoint]
    else:
        return (x[midpoint] + (x[midpoint - 1]))/2.0

def IQR(x):
    """calculates the interquartile range of list x
    
    returns min_val, max_val"""
    
    x.sort()
    #split values into lower and upper portions at the median
    odd = len(x) % 2
    midpoint = int(len(x)/2)
    if odd:
        low_vals = x[:midpoint]
        high_vals = x[midpoint+1:]
    else: #if even
        low_vals = x[:midpoint]
        high_vals = x[midpoint:]
    #find the median of the low and high values
    min_val = median(low_vals)
    max_val = median(high_vals)
    return min_val, max_val

class Jackknifer(object):
    """Abstract class that allows for jackknifing.
    
    This is an abstract class, it should only be instantiated via its
    subclasses
    
    Specs for sub-classing:
        The programmer must define four functions:
            _remove_items()
                knows how to remove the appropriate number of items for the
                type of object being jackknifed
            _evaluate_true()
                returns a true value for the object before any items are 
                removed
            _evaluate_reduced()
                returns a list of values for reduced objects
            _compare_values() 
                specifies how the true and reduced values should be compared
        """
        
    def __init__(self, obj, num_reps, value, remove_by, keep):
        """Returns new jackknifer, an abstract class

        obj: the object to be jackknifed
        num_reps: the number of times to subsample the data
        value: represent the number/fraction of items to remove/keep
        remove_by: if == fraction, treats value as a fraction
                    if == number, treats value as a number
        keep: if True, determined items to keep based on value
                if False, removes items based on value"""

        self.obj = obj
        self.num_reps = num_reps
        self.value = value
        self.remove_by = remove_by
        self.keep = keep

        self.TrueValue = None
        self.jackknifedValues = None
        
    def _number_to_remove(self, size):
        """Determines number of descendants to remove from the tree"""
        if self.remove_by == 'fraction':
            result = round(self.value * size)
            if self.keep:
                result = size - result
        elif self.remove_by == 'number':
            result = self.value
            if self.keep:
                result = size - self.value
        else:
            return ValueError, "remove_by must be set to 'number' or 'fraction'"
        return result
    
    def __call__(self):
        """Performs jackknifing"""
        self._evaluate_true()
        self._evaluate_reduced()
        return self._compare_values()
       
class JackknifeTree(Jackknifer):
    """Abstract class for Tree jackknifing classes
    
    tree is an EnvsNode object"""

    def __init__(self, obj, num_reps, value, remove_by='number', keep=False, \
            Weight=False, norm=False):
        """adds Weight to SuperClass __init__"""
        #Run __init__ from parent class
        Jackknifer.__init__(self, obj, num_reps, value, remove_by, keep)
        self.Weight = Weight
        self.norm = norm

class JackknifeTreeEnvSpecies(JackknifeTree):
    """Class for Jackknifing trees by removing species from a particular env"""

    def __init__(self, obj, num_reps, value, remove_by='number', keep=False, \
            Weight=False, env=None):
        """Adds env to superclass __init__"""
        JackknifeTree.__init__\
                (self, obj, num_reps, value, remove_by, keep, Weight)
        self.env = env
        
    def _remove_items(self, number_to_remove):
        """Removes number_to_remove of leaves in env from a copy of self.obj.
        """
        obj_copy = deepcopy(self.obj)
        # must assign these because program updates them when removing items
        obj_copy.assignInternalEnvs()
        env_term_decs = obj_copy.returnEnvTermDecs(self.env)
        if len(env_term_decs) >= number_to_remove:
            obj_copy.removeEnvTermNodes(self.env, number_to_remove)
            return obj_copy
        else:
            raise ValueError, "cannot remove more objects than are in tree."

class JackknifeTreeSpeciesUPGMA(JackknifeTreeEnvSpecies):
    """Class for jackknifing  UPGMA cluster to removed species from an env"""

    def __init__(self, obj, num_reps, value, remove_by='number', keep=False, \
            Weight=False, env=None):
        """Run JackknifeTreeEnvSpecies __init__"""
        JackknifeTreeEnvSpecies.__init__(self, obj, num_reps, value, \
                remove_by, keep, Weight, env)

    def _evaluate_true(self):
        """Returns the UPGMA cluster of the true tree."""
        cluster = self.obj.clusterEnvsUPGMA(self.Weight)
        cluster.nameUnnamedNodes()
        self.TrueValue, node_list = cluster.makeTreeArray()

    def _evaluate_reduced(self):
        """Produces a UPGMA cluster for num_reps reduced trees"""
        size = len(self.obj.TerminalDescendants)
        remove = self._number_to_remove(size)
        results = []
        for rep in range(self.num_reps):
            reduced = self._remove_items(remove)
            cluster = reduced.clusterEnvsUPGMA(self.Weight)
            cluster.nameUnnamedNodes()
            clusterArray, node_list = cluster.makeTreeArray()
            results.append(clusterArray)
        self.jackknifedValues = results
       
    def _compare_values(self):
        """Compares true UPGMA cluster to reduced UPGMA clusters"""

        comparisons = \
                [1 for val in self.jackknifedValues if val != self.TrueValue]
        return 1 - (sum(comparisons)/self.num_reps)

class JackknifeEnvs(Jackknifer):
    """Abstract class for Environment jackknifing classes"""
    
    def __init__(self, obj, num_reps, value, remove_by='number', keep=False, \
            Weight=False):
        """add Weight and dist_matrix to superclass __init__"""
        Jackknifer.__init__(self, obj, num_reps, value, remove_by, keep)
        self.Weight = Weight
        self.dist_matrix = self.obj.makeEnvMatrix(BIG_NUM, Weight)
    
    def _remove_items(self, number_to_remove):
        """Removes number_to_remove of envs from an dist_matrix for a tree"""
        reduced_dist_matrix = deepcopy(self.dist_matrix)
        indices = []
        for i in range(number_to_remove):
            #select a random env index to remove and append to indices
            index = randrange(0, len(self.dist_matrix), 1)
            indices.append(index)
            #remove the corresponding row and column from the dist_matrix
            del reduced_dist_matrix.RowOrder[index]
            reduced_dist_matrix.purge()
        return reduced_dist_matrix, indices

class JackknifeEnvsUPGMA(JackknifeEnvs):
    """Class for measuring robustness of UPGMA clusters to included Envs"""
    
    def __init__(self, obj, num_reps, value, remove_by='number', keep=False, \
            Weight=False):
        #run __init__ from parent class
        JackknifeEnvs.__init__(self, obj, num_reps, value, remove_by, keep, \
                Weight)
        self.node_list = None
        self.cluster = None

    def _evaluate_true(self):
        """Sets self.TrueValue to an env_array for a UPGMA cluster."""
        #cluster with UPGMA and make a bit representation of the cluster
        dist_matrix_array, node_order = inputs_from_dict2D(self.dist_matrix)
        self.cluster = UPGMA_cluster(dist_matrix_array, node_order, BIG_NUM)
        self.cluster.nameUnnamedNodes()
        self.TrueValue, self.node_list = self.cluster.makeTreeArray()

    def _evaluate_reduced(self):
        """Sets self.jackknifedValues to a list of arrays for reduced trees."""
        size = len(self.dist_matrix)
        remove = self._number_to_remove(size)
        results = []
        for rep in range(self.num_reps):
            #remove the specified numnber of envs from the distance matrix
            reduced_dist_matrix, indices = self._remove_items(remove)
            #make a new UPGMA cluster based on the modified tree
            reduced_dist_matrix_array, reduced_node_list = \
                    inputs_from_dict2D(reduced_dist_matrix)
            reduced_cluster = UPGMA_cluster(reduced_dist_matrix_array, \
                    reduced_node_list, BIG_NUM)
            reduced_cluster.nameUnnamedNodes()
            #use dec_list from initial cluster to make clusterArray so that
            #old and new can be compared. clusterArray2 will have a column
            #for the missing env but its value will be 0 for all nodes
            dec_list = [dec.Data for dec in self.cluster.TerminalDescendants]
            dec_list.sort()
            reduced_clusterArray, x = reduced_cluster.makeTreeArray(dec_list)
            results.append([reduced_clusterArray, indices])
        self.jackknifedValues = results

    def _compare_values(self):
        """Returns a dict with a jackknife value for each node in self.cluster."""
        #initialize results dictionary
        result = {}
        for node in self.node_list:
            result[node.Data] = []
        #compare self.TrueValue to each array in self.jackknifedValues
        #append a 1 to the result if a node in the original tree is in 
        #the modified tree and a 0 if it is not
        for reduced_tree_info in self.jackknifedValues:
            reduced_clusterArray = reduced_tree_info[0]
            indices = reduced_tree_info[1]
            #set the column in the original array to 0 for each removed env
            #in the reduced array
            TrueValue_mod = deepcopy(self.TrueValue)
            for index in indices:
                TrueValue_mod[:, index] = 0
            for i_r, row in enumerate(TrueValue_mod):
                in_array = 0
                #if the env removed results in a node with only 1 dec
                #record the node as still being there
                if sum(row) == 1:
                    in_array = 1
                #also set to 1 if the same node exists in the new tree
                for row2 in reduced_clusterArray:
                    if list(row) == list(row2):
                        in_array = 1
                result[self.node_list[i_r].Data].append(in_array)
        #replace the list values in result with jackknife fractions
        for i in result:
            result[i] = sum(result[i])/len(result[i])
        return result


class JackknifeSpeciesBalancedPCA(JackknifeTree):
    """Class for measuring robustness of PCA clusters to # seqs
    
    Makes a balanced tree and draws confidence ellipses
    around PCA plots

    span:
        95CI: uses 95% confidence intervals
        99CI: uses 99% confidence intervals
        min_max: uses minimum and maximum value
        IQR: inter-quartile range
    """   
    
    def __init__(self, obj, num_reps, value, remove_by='number', keep=False, \
            Weight=False, norm=False, span='95CI'):
        #run __init__ from parent class
        JackknifeTree.__init__(self, obj, num_reps, value, remove_by, keep, \
                Weight, norm)
        self.span = span

    def _evaluate_true(self):
        """Returns the PCA results of the true tree."""
        true_tree = deepcopy(self.obj)
        dis_mat = true_tree.makeEnvMatrix(0.0, self.Weight, self.norm)
        matrix, node_order = inputs_from_dict2D(dis_mat)
        names = [node.Data for node in node_order]
        point_matrix, eigvals = principal_coordinates_analysis(matrix)
        self.TrueValue = (point_matrix, eigvals, names)
    
    def _evaluate_reduced(self):
        """Performs PCA on num_reps balanced trees"""
        results = []
        for rep in range(self.num_reps):
            jn_tree = deepcopy(self.obj)
            jn_tree.balanceTreeEnvs(self.value, self.Weight)
            dis_mat = jn_tree.makeEnvMatrix(0.0, self.Weight, self.norm)
            matrix, node_order = inputs_from_dict2D(dis_mat)
            names = [node.Data for node in node_order]
            point_matrix, eigvals = principal_coordinates_analysis(matrix)
            results.append((point_matrix, eigvals, names))
        self.jackknifedValues = results

    def _compare_values(self):
        """returns average and 95% confidence interval for PCA results"""
        #make a dict2D containing lists of jackknife values for each PC value
        jn_results = self._compile_jn_results()
        f = open('jn_results.pi', 'w')
        pickle.dump(jn_results, f)
        return self._compute_results(jn_results)

    def _flip_vectors(self, jn_matrix, jn_eigvals):
        """transforms PCA vectors so that signs are correct"""
        new_matrix= zeros([len(jn_matrix), len(jn_matrix[0])], Float)
        #get info for the True PCA matrix
        true_matrix = self.TrueValue[0]
        true_eigvals = self.TrueValue[1]
        true_vector_order = list(argsort(true_eigvals))
        true_vector_order.reverse()
        #get info for the jn PCA matrix
        jn_vector_order = list(argsort(jn_eigvals))
        jn_vector_order.reverse()
        for pc_n, pc_i in enumerate(jn_vector_order):
            true_vector = true_matrix[true_vector_order[pc_n]]
            jn_vector = jn_matrix[pc_i]
            disT = list(true_vector - jn_vector)
            disT = sum(map(abs, disT))
            jn_flip = jn_vector*[-1]
            disF = list(true_vector - jn_flip)
            disF = sum(map(abs, disF))
            if disT > disF:
                new_matrix[pc_i] = jn_flip
            else:
                new_matrix[pc_i] = jn_vector
        return new_matrix
            
    def _compile_jn_results(self):
        """compiles results of all jn runs into 1 dict2D
        
        compare each vector to true vector and flip signs if needed"""
        jn_results = {}
        for jn_rep in self.jackknifedValues:
            matrix = jn_rep[0]
            eigvals = jn_rep[1]
            matrix = self._flip_vectors(matrix, eigvals)
            names = jn_rep[2]
            vector_order = list(argsort(eigvals))
            vector_order.reverse()
            for name_i, name in enumerate(names):
                if name not in jn_results:
                    jn_results[name] = {}
                for pc_n, pc_i in enumerate(vector_order):
                    if str(pc_n+1) not in jn_results[name]:
                        jn_results[name][str(pc_n+1)] = [matrix[pc_i, name_i]]
                    else:
                        jn_results[name][str(pc_n+1)].append(matrix[pc_i, name_i])
        return Dict2D(jn_results)
                    
    def _compute_results(self, jn_results):
        """returns averages and ranges (based on 95% CI) for results"""
        average_values = []
        ranges = []
        for env1 in jn_results:
            for PC in jn_results[env1]:
                values = Numbers(jn_results[env1][PC])
                if self.span == '95CI':
                    SE = values.StandardDeviation/sqrt(len(values))
                    r = 2 * 1.96 * SE
                    midpoint = values.Mean
                if self.span == '99CI':
                    SE = values.StandardDeviation/sqrt(len(values))
                    r = 2 * 2.58 * SE
                    midpoint = values.Mean
                elif self.span == 'min_max':
                    r = max(values) - min(values)
                    midpoint = max(values) - (r/2.0)
                elif self.span == 'IQR':
                    min_val, max_val = IQR(values)
                    r = max_val - min_val
                    midpoint = max_val - (r/2)
                average_values.append((env1, PC, midpoint))
                ranges.append((env1, PC, r))
        return Dict2D(average_values), Dict2D(ranges)


class JackknifeSpeciesBalancedPCACorr(JackknifeSpeciesBalancedPCA):
    """Class for measuring robustness of PCA correlations to # seqs
    
    Makes a balanced tree and measures average and standard dev of
    the correlation of input PC number with an environmental variable

    env_var is a dict mapping env_name to a measurement (e.g. temperature)
    PC_num is the principal coordinate number to check for the correlation with
    """   
    
    def __init__(self, obj, num_reps, value, remove_by='number', keep=False, \
            Weight=False, norm=False, env_var={}, PC_num = 1):
        #run __init__ from parent class
        JackknifeTree.__init__(self, obj, num_reps, value, remove_by, keep, \
                Weight, norm)
        self.env_var = env_var
        self.PC_num = PC_num

    def _compare_values(self):
        "returns average and std dev of correlations of PC vector with env_var"
        corr_values = []
        names = self.jackknifedValues[0][2]
        print "names=", names
        env_var_values = []
        prob_values = []
        for name in names:
            env_var_values.append(self.env_var[name])
        print "env_var=", env_var_values
        for jn_rep in self.jackknifedValues:
            matrix = jn_rep[0]
            eigvals = jn_rep[1]
            matrix = self._flip_vectors(matrix, eigvals)
            vector_order = list(argsort(eigvals))
            vector_order.reverse()
            pc_vec = matrix[vector_order[self.PC_num-1]]
            print "pc_vec=", pc_vec
            r, prob = correlation(pc_vec, env_var_values)
            corr_values.append(r)
            prob_values.append(prob)
            print "corr_values=", corr_values
            print "prob =", prob
            print "***********************"
        corr_values = Numbers(corr_values)
        return corr_values.Mean, corr_values.StandardDeviation
        
